from libc.stdlib cimport malloc, free
from libcpp.string cimport string
from libcpp.vector cimport vector
from cython.operator cimport dereference as deref, preincrement as inc
from cython import sizeof, NULL

import socket

import bson

from .collection_space import CollectionSpace
from .collection import Collection
from .error import Error

cdef extern from "ossErr.h":
    enum:SDB_OK
    enum:SDB_DMS_EOC

cdef extern from "ossTypes.h":
    ctypedef char CHAR

    ctypedef unsigned long long UINT64
    ctypedef long long SINT64
    ctypedef long long INT64

    ctypedef int INT32
    ctypedef int SINT32
    ctypedef unsigned int UINT32

    ctypedef unsigned short UINT16
    ctypedef short SINT16
    ctypedef short INT16

    ctypedef INT32 BOOLEAN

    INT32 ossRoundUpToMultipleX(INT32, INT32)

cdef extern from "ossVer.h":
    void ossGetVersion(INT32*, int*, int*, const char**)

cdef extern from "msg.h":
    ctypedef struct MsgSysInfoRequest:
        pass
    ctypedef struct MsgSysInfoReply:
        pass
    ctypedef struct MsgOpReply:
        pass

cdef extern from "common.h":

    INT32 clientBuildSysInfoRequest ( CHAR **, INT32 * )
    INT32 clientExtractSysInfoReply ( CHAR *, BOOLEAN *, INT32 * )

    INT32 md5Encrypt( const CHAR *, CHAR*, UINT32)


    INT32 clientBuildAuthMsg( CHAR **ppBuffer, INT32 *bufferSize,
                          const CHAR *pUsrName,
                          const CHAR *pPasswd,
                          UINT64 reqID, BOOLEAN endianConvert ) ;

    INT32 clientExtractReply ( CHAR *pBuffer, SINT32 *flag, SINT64 *contextID,
                           SINT32 *startFrom, SINT32 *numReturned,
                           BOOLEAN endianConvert )

cdef extern from "_sdbapi.h":
    enum:MD5_DIGEST_LENGTH

cdef extern from "_sdbapi.h" namespace "qc":
    SINT32 steal_size(const CHAR*)
    ctypedef void (*network_func)(CHAR* buff, size_t size)

    cdef cppclass ConnectionContext:
        ConnectionContext() except +

        CHAR* send_buff
        INT32 send_buff_size
        CHAR* recv_buff
        INT32 recv_buff_size
        BOOLEAN endian_convert
        network_func sender
        network_func recver

        string str()

    INT32 build_create_collection_space_command(
        ConnectionContext* conn_ctx,
        const CHAR* cs_name,
        INT32 page_size,
    )

    INT32 build_drop_collection_space_command(
        ConnectionContext* conn_ctx,
        const CHAR* cs_name
    )

    INT32 build_create_collection_command(
        ConnectionContext* conn_ctx,
        const CHAR* cs_name,
        const CHAR* c_name,
        const CHAR* options
    )

    INT32 build_drop_collection_command(
        ConnectionContext* conn_ctx,
        const CHAR* cs_name,
        const CHAR* c_name
    )

    INT32 build_insert_command(
        ConnectionContext* conn_ctx,
        const CHAR* cs_name,
        const CHAR* c_name,
        const CHAR* bytes_record
    )

    INT32 build_bulk_insert_command(
        ConnectionContext* conn_ctx,
        const CHAR* cs_name,
        const CHAR* c_name,
        SINT32 flag,
        vector[string]& bytes_record
    )

    INT32 build_cr_next_command(
        ConnectionContext* conn_ctx,
        SINT64 context_id,
        INT64 total_read
    )

    INT32 build_query_command(
        ConnectionContext* conn_ctx,
        const CHAR* cs_name,
        const CHAR* c_name,
        const CHAR* bytes_condition,
        const CHAR* bytes_selector,
        const CHAR* bytes_order_by,
        const CHAR* bytes_hint,
        INT64 num_to_skip,
        INT64 num_to_return,
        INT32 flag
    )

cdef class Connection(object):
    cdef ConnectionContext *_ctx
    cdef object _sock
    def __cinit__(self):
        self._ctx = new ConnectionContext()

    def __dealloc__(self):
        del self._ctx

    def connect(self, address):
        try:
            if self._sock:
                raise Error(
                    "Socket is connected, cann't connnect to %s" % address)
        except AttributeError:
            pass
        self._sock = socket.create_connection(address)
        self._do_connect()

    def close(self):
        ret = self._sock.close()
        self._sock = None
        return ret

    def _do_connect(self):
#        print self._ctx.str()
#        print '*'*30
        cdef INT32 rc = clientBuildSysInfoRequest(
            &(self._ctx.send_buff),
            &(self._ctx.send_buff_size))
        if rc:
            raise Error("clientBuildSysInfoRequest failed.")
#        print self._ctx.str()
#        print '*'*30
        self.send(sizeof(MsgSysInfoRequest))
        pybuff = self.recv(sizeof(MsgSysInfoReply))
        cdef CHAR* buff = pybuff
        rc = clientExtractSysInfoReply(
            buff,
            &(self._ctx.endian_convert),
            NULL);
        if rc:
            raise Error("clientExtractSysInfoReply failed.")
#        print self._ctx.str()
#        print '*'*30

    def send(self, size_t size):
#        cdef size_t sent = 0
#        while sent < size:
#            sent += self._sock.send(self._ctx.send_buff[sent:size])
        self._sock.sendall(self._ctx.send_buff[:size])
#        print 'SEND:', size, repr(self._ctx.send_buff[:size])

    def recv(self, size_t size):
        buff = self._sock.recv(size, socket.MSG_WAITALL)
#        while len(buff) < size:
#            buff += self._sock.recv(size - len(buff), socket.MSG_WAITALL)
#        print 'RECV:', size, repr(buff)
        return buff

    cdef peek(self, size_t size):
        while True:
            buff = self._sock.recv(size, socket.MSG_PEEK | socket.MSG_WAITALL)
            if len(buff) < size:
                continue
#            print 'PEEK:', size, repr(buff)
            return buff

    cdef SINT32 _peek_msg_len(self):
        p = self.peek(sizeof(SINT32))
        cdef CHAR* buff = p
        return deref(<SINT32*>buff)

    
    def check_user(self, const CHAR* username, const CHAR* password):
        cdef CHAR md5[MD5_DIGEST_LENGTH]
        cdef INT32 rc = md5Encrypt(password, md5, MD5_DIGEST_LENGTH)
        if rc:
            raise Error("md5Encrypt failed.")
        rc = clientBuildAuthMsg(
            &(self._ctx.send_buff),
            &(self._ctx.send_buff_size),
            username,
            md5,
            0,
            self._ctx.endian_convert)
        if rc:
            raise Error("clientBuildAuthMsg failed.")
        self.send(steal_size(self._ctx.send_buff))
        # recv
        cdef SINT32 length = self._peek_msg_len()
        buff = self.recv(length)
        cdef BOOLEAN reply_flag
        cdef SINT64 context_id
        cdef INT32 start_from
        cdef INT32 num_returned
        rc = clientExtractReply(<CHAR*>buff,
            &reply_flag,
            &context_id,
            &start_from,
            &num_returned,
            self._ctx.endian_convert)
        if rc:
            raise Error("clientExtractReply failed.")
        if reply_flag != SDB_OK:
            self.close()
            raise Error("password error.")
    
    def create_collection_space(self, CHAR* cs_name, INT32 page_size):
        cdef INT32 rc = build_create_collection_space_command(
            self._ctx,
            cs_name,
            page_size)
        if SDB_OK != rc:
            raise Error("build_create_collection_space_command failed.")
        self.send(steal_size(self._ctx.send_buff))
        # recv
        cdef SINT32 length = self._peek_msg_len()
        buff = self.recv(length)
        cdef BOOLEAN reply_flag
        cdef SINT64 context_id
        cdef INT32 start_from
        cdef INT32 num_returned
        rc = clientExtractReply(<CHAR*>buff,
            &reply_flag,
            &context_id,
            &start_from,
            &num_returned,
            self._ctx.endian_convert)
        if rc:
            raise Error("clientExtractReply failed.")
        if reply_flag != SDB_OK:
            self.close()
            raise Error("create collection space failed.")

        return CollectionSpace(self, cs_name)

    def drop_collection_space(self, CHAR* cs_name):
        cdef INT32 rc = build_drop_collection_space_command(
            self._ctx,
            cs_name)
        if SDB_OK != rc:
            raise Error("build_drop_collection_space_command failed.")
        self.send(steal_size(self._ctx.send_buff))
        # recv
        cdef SINT32 length = self._peek_msg_len()
        buff = self.recv(length)
        cdef BOOLEAN reply_flag
        cdef SINT64 context_id
        cdef INT32 start_from
        cdef INT32 num_returned
        rc = clientExtractReply(<CHAR*>buff,
            &reply_flag,
            &context_id,
            &start_from,
            &num_returned,
            self._ctx.endian_convert)
        if rc:
            raise Error("clientExtractReply failed.")
        if reply_flag != SDB_OK:
            self.close()
            raise Error("drop collection space failed.")

    def create_collection(self, cs_name, c_name, options):
        cdef INT32 rc = build_create_collection_command(
            self._ctx,
            cs_name,
            c_name,
            options)
        if SDB_OK != rc:
            raise Error("build_create_collection_command failed.")
        self.send(steal_size(self._ctx.send_buff))
        # recv
        cdef SINT32 length = self._peek_msg_len()
        buff = self.recv(length)
        cdef BOOLEAN reply_flag
        cdef SINT64 context_id
        cdef INT32 start_from
        cdef INT32 num_returned
        rc = clientExtractReply(<CHAR*>buff,
            &reply_flag,
            &context_id,
            &start_from,
            &num_returned,
            self._ctx.endian_convert)
        if rc:
            raise Error("clientExtractReply failed.")
        if reply_flag != SDB_OK:
            self.close()
            raise Error("create collection failed.")

        return Collection(self, cs_name, c_name)

    def drop_collection(self, cs_name, c_name):
        cdef INT32 rc = build_drop_collection_command(
            self._ctx,
            cs_name,
            c_name)
        if SDB_OK != rc:
            raise Error("build_drop_collection_command failed.")
        self.send(steal_size(self._ctx.send_buff))
        # recv
        cdef SINT32 length = self._peek_msg_len()
        buff = self.recv(length)
        cdef BOOLEAN reply_flag
        cdef SINT64 context_id
        cdef INT32 start_from
        cdef INT32 num_returned
        rc = clientExtractReply(<CHAR*>buff,
            &reply_flag,
            &context_id,
            &start_from,
            &num_returned,
            self._ctx.endian_convert)
        if rc:
            raise Error("clientExtractReply failed.")
        if reply_flag != SDB_OK:
            self.close()
            raise Error("drop collection failed.")

    def insert(self, cs_name, c_name, bytes_record):
        cdef INT32 rc = build_insert_command(
            self._ctx,
            cs_name,
            c_name,
            bytes_record)
        if SDB_OK != rc:
            raise Error("build_insert_command failed.")
        self.send(steal_size(self._ctx.send_buff))
        # recv
        cdef SINT32 length = self._peek_msg_len()
        buff = self.recv(length)
        cdef BOOLEAN reply_flag
        cdef SINT64 context_id
        cdef INT32 start_from
        cdef INT32 num_returned
        rc = clientExtractReply(<CHAR*>buff,
            &reply_flag,
            &context_id,
            &start_from,
            &num_returned,
            self._ctx.endian_convert)
        if rc:
            raise Error("clientExtractReply failed.")
        if reply_flag != SDB_OK:
            self.close()
            raise Error("insert failed.")

    def bulk_insert(self, cs_name, c_name, flag, bytes_records):
        cdef vector[string] records
        for br in bytes_records:
            records.push_back(<bytes>br)

        cdef INT32 rc = build_bulk_insert_command(
            self._ctx,
            cs_name,
            c_name,
            flag,
            records)
        if SDB_OK != rc:
            raise Error("build_insert_command failed.")
        self.send(steal_size(self._ctx.send_buff))
        # recv
        cdef SINT32 length = self._peek_msg_len()
        buff = self.recv(length)
        cdef BOOLEAN reply_flag
        cdef SINT64 context_id
        cdef INT32 start_from
        cdef INT32 num_returned
        rc = clientExtractReply(<CHAR*>buff,
            &reply_flag,
            &context_id,
            &start_from,
            &num_returned,
            self._ctx.endian_convert)
        if rc:
            raise Error("clientExtractReply failed.")
        if reply_flag != SDB_OK:
            self.close()
            raise Error("insert failed.")

    def query(self, cs_name, c_name, bson_condition, bson_selector,
            bson_order_by, bson_hint, num_to_skip, num_to_return):
        cdef INT32 flag = 0
        cdef CHAR* bytes_condition
        cdef CHAR* bytes_selector
        cdef CHAR* bytes_order_by
        cdef CHAR* bytes_hint
        if bson_condition:
            bytes_condition = <CHAR*>bson_condition
        else:
            bytes_condition = NULL
        if bson_selector:
            bytes_selector = <CHAR*>bson_selector
        else:
            bytes_selector = NULL
        if bson_order_by:
            bytes_order_by = <CHAR*>bson_order_by
        else:
            bytes_order_by = NULL
        if bson_hint:
            bytes_hint = <CHAR*>bson_hint
        else:
            bytes_hint = NULL
        cdef INT32 rc = build_query_command(
            self._ctx,
            cs_name,
            c_name,
            bytes_condition,
            bytes_selector,
            bytes_order_by,
            bytes_hint,
            num_to_skip,
            num_to_return,
            flag)
        if SDB_OK != rc:
            raise Error('build_query_command failed.')

        self.send(steal_size(self._ctx.send_buff))
        # recv
        cdef SINT32 length = self._peek_msg_len()
        buff = self.recv(length)
        cdef BOOLEAN reply_flag
        cdef SINT64 context_id
        cdef INT32 start_from
        cdef INT32 num_returned
        rc = clientExtractReply(<CHAR*>buff,
            &reply_flag,
            &context_id,
            &start_from,
            &num_returned,
            self._ctx.endian_convert)
        if rc:
            raise Error("clientExtractReply failed.")
        if reply_flag != SDB_OK:
            self.close()
            raise Error("query failed.")
        
        return Cursor(self, context_id)
    
    def cr_next(self, SINT64 context_id, INT64 total_read):
        cdef INT32 rc = build_cr_next_command(
            self._ctx,
            context_id,
            total_read)
        if SDB_OK != rc:
            raise Error("build_cr_next_command failed.")
        self.send(steal_size(self._ctx.send_buff))
        # recv
        cdef SINT32 length = self._peek_msg_len()
        buff = self.recv(length)
        cdef BOOLEAN reply_flag
        cdef SINT64 cid
        cdef INT32 start_from
        cdef INT32 num_returned
        rc = clientExtractReply(<CHAR*>buff,
            &reply_flag,
            &cid,
            &start_from,
            &num_returned,
            self._ctx.endian_convert)
        if rc:
            raise Error("clientExtractReply failed.")
        if cid != context_id:
            raise Error("context_id error.")
        if reply_flag == SDB_DMS_EOC:
            raise StopIteration
        elif reply_flag != SDB_OK:
            self.close()
            raise Error("cr_next failed.")
#        cdef MsgOpReply* reply = <MsgOpReply*>(<CHAR*>buff);
        cdef INT32 offset = ossRoundUpToMultipleX(sizeof(MsgOpReply), 4)
        cdef INT32 size = steal_size(buff)
        return buff[offset:size]

cdef class Cursor(object):
    cdef INT64 _total_read
    cdef SINT64 _context_id
    cdef object _conn
    cdef bytes _buff
    cdef INT32 _offset
    def __cinit__(self, conn, SINT64 context_id):
        self._conn = conn
        self._context_id = context_id
        self._total_read = 0L
        self._buff = b''
        self._offset = -1

    def next(self):
        if len(self._buff) == 0 \
                or self._offset <= -1 \
                or self._offset >= len(self._buff):
            self._buff = self._conn.cr_next(self._context_id, self._total_read)
            self._offset = 0
        bytes_record = self._read_from_buff()
        inc(self._total_read)
        return self._convert_to_dict(bytes_record)

    cdef bytes _read_from_buff(self):
        cdef INT32 bson_size = steal_size(self._buff[self._offset:])
        cdef bytes buff = self._buff[self._offset:self._offset+bson_size]
        cdef INT32 packet_size = ossRoundUpToMultipleX(bson_size, 4)
#        print '_offset', self._offset, 'bson-size', bson_size, 'packet-size', packet_size, 'buff', repr(buff)
        self._offset += packet_size
        return buff

    def current(self):
        bytes_record = self._conn.cr_current(self._context_id, self._total_read)
        return self._convert_to_dict(bytes_record)

    def _convert_to_dict(self, bytes_record):
        record, size = bson._bson_to_dict(
            bytes_record,
            dict,
            False,
            bson.OLD_UUID_SUBTYPE,
            True)
        return record

def get_version():
    cdef int version, subversion, release
    cdef char* build = NULL
    cdef bytes pys_build
    ossGetVersion(&version, &subversion, &release, <const char**>&build)
    pys_build = build
    return version, subversion, release, pys_build

def create_collection_space(conn, cs_name, page_size):
    return conn.create_collection_space(cs_name, page_size)


